use std::collections::HashMap;

use crate::mir::{self, basicblock, function, instr::{self, Load, Operand}, mir_ref::{BasicBlockRef, FunctionRef, InstrRef, MIRSlotMap}, mtype, reg::Vreg};

pub struct Context{
    pub bb_cnt: usize,
    pub instr_cnt: usize,
    // pub cur_bb: Option<BasicBlockRef>,
    pub verg_cnt: usize,
}

impl Context {
    pub fn new() -> Self {
        Self {
            bb_cnt : 0,
            instr_cnt : 1,
            // cur_bb : None,
            verg_cnt: 1,
        }
    }

    pub fn get_new_bb_idx(&mut self) -> usize {
        let old = self.bb_cnt.clone();
        self.bb_cnt += 1;
        return old;
    }

    pub fn get_new_instr_id(&mut self) -> usize {
        let old = self.instr_cnt.clone();
        self.instr_cnt += 1;
        return old;
    }

    pub fn get_new_verg_id(&mut self) -> usize {
        let old = self.verg_cnt.clone();
        self.verg_cnt += 1;
        old
    }
}

pub struct MirBuilder<'a> {
    pub ctx: &'a mut Context,
    pub cur_func: Option<FunctionRef>,
    pub cur_bb: Option<BasicBlockRef>,
    pub mir_arena: MIRSlotMap,
    pub instr2reg: HashMap<InstrRef, Vreg>,
}

impl<'a> MirBuilder<'a> {
    pub fn new(ctx: &'a mut Context) -> Self {
        Self { 
            ctx: ctx,
            cur_func: None,
            cur_bb: None,
            mir_arena: MIRSlotMap::new(),
            instr2reg: HashMap::new(),            
        }
    }

    pub fn create_func(&mut self, 
        func_name: String,
        func_ret_ty: mtype::MIRType,
        func_f_params: Vec<(String, mtype::MIRType)>) -> FunctionRef 
    {

        let mut new_func = function::func::new(func_name, func_ret_ty, func_f_params);
        let new_bb = basicblock::bb::new(self.ctx.get_new_bb_idx());

        let new_bb_ref = self.mir_arena.add_bb(new_bb);
        self.cur_bb = Some(new_bb_ref.clone());

        new_func.add_bb(new_bb_ref);
    
        let fr = self.mir_arena.add_func(new_func);
        self.set_cur_func(fr);

        fr
    }

    pub fn set_cur_func(&mut self, f: FunctionRef) {
        self.cur_func = Some(f);
    }

    pub fn create_new_bb(&mut self) -> BasicBlockRef {
        let new_bb = basicblock::bb::new(self.ctx.get_new_bb_idx());
        let br = self.mir_arena.add_bb(new_bb);
        let cur_func = self.mir_arena.get_mut_func(self.cur_func.unwrap()).unwrap();
        cur_func.bbs.push(br);
        br
    }

    pub fn set_cur_bb(&mut self, bb: BasicBlockRef) {
        self.cur_bb = Some(bb);
    }

    fn add_instr2cur_bb(&mut self, i: InstrRef) {
        match self.cur_bb {
            Some(bb_ref) => {
                let mut bb = self.mir_arena.get_mut_bb(bb_ref).unwrap();
                bb.instrs.push(i);
            }
            None => {}
        }
    }

    pub fn get_instr_vreg(&mut self, instr: &InstrRef) -> Vreg {
        let r = self.instr2reg.get(&instr);

        // 已经有虚拟寄存器，表明指令是已经发射了
        match r {
            Some(reg) => reg.clone(),
            None => {
                let new_vreg = Vreg::new(self.ctx.get_new_verg_id());
                let minstr = self.mir_arena.get_mut_instr(*instr).unwrap();
                minstr.vreg_id = Some(new_vreg.clone());
                // 获取虚拟寄存器号后， 发射指令
                match self.cur_bb {
                    Some(bb) => {
                        self.add_instr2cur_bb(*instr);
                    }
                    None => {}
                }
                new_vreg
            }
        }
    }

    pub fn create_bin(&mut self, binOp: instr::BinOp ,lhs: Operand , rhs: Operand) -> InstrRef {
        // 这里使用到了lhs，rhs，要检查值是否在寄存器上，指令也是先放在池里，等要用到了给分配一个虚拟寄存器，加载进去
        match lhs {
            instr::Operand::ConstValue(_) => {}
            instr::Operand::Vreg(v) => {
                let _ = self.get_instr_vreg(&v);
            }
        };

        match rhs {
            instr::Operand::ConstValue(_) => {}
            instr::Operand::Vreg(v) => {
                let _ = self.get_instr_vreg(&v);
            }
        };

        let bin_instr = instr::BinInstr::new(binOp, lhs, rhs);
        let m_instr = instr::MirInst::new(
            instr::MirInstKind::BinInst(bin_instr), 
            self.ctx.get_new_instr_id(), None);

        let res = self.mir_arena.add_instr(m_instr);

        // self.add_instr2cur_bb(res);
        res
    }

    pub fn create_return(&mut self, ret_val: Option<Operand>) -> InstrRef {

        let ret_instr = instr::ReturnInstr::new(ret_val);
        let m_instr = instr::MirInst::new(
            instr::MirInstKind::ReturnInstr(ret_instr), 
            self.ctx.get_new_bb_idx(), None);

        let res = self.mir_arena.add_instr(m_instr);

        self.add_instr2cur_bb(res);
        res
    }

    pub fn create_switch(&mut self, 
        cond: Operand, 
        true_bb: BasicBlockRef, 
        false_bb: Option<BasicBlockRef>) -> InstrRef 
    {
        match cond {
            instr::Operand::ConstValue(_) => {}
            instr::Operand::Vreg(v) => {
                let _ = self.get_instr_vreg(&v);
            }
        };

        let si = mir::instr::SwitchInstr::new(cond, true_bb, false_bb);
        let instr = mir::instr::MirInst::new(
            mir::instr::MirInstKind::SwitchInstr(si), 
            self.ctx.get_new_instr_id(), 
            None);

        let i = self.mir_arena.add_instr(instr);
        self.add_instr2cur_bb(i);
        i
    }

    pub fn create_goto(&mut self, dst: BasicBlockRef) -> InstrRef {
        let goto_instr = mir::instr::GotoInstr::new(dst);
        let instr = mir::instr::MirInst::new(
            mir::instr::MirInstKind::GotoInstr(goto_instr), 
            self.ctx.get_new_instr_id(), 
            None);
        let i = self.mir_arena.add_instr(instr);
        self.add_instr2cur_bb(i);
        i
    }

    pub fn create_load(&mut self, dst: String) -> InstrRef {
        // load to vreg
        let ld = mir::instr::Load::new(dst);
        let vreg = mir::reg::Vreg::new(self.ctx.get_new_verg_id());
        let instr = mir::instr::MirInst::new(
            mir::instr::MirInstKind::Load(ld), 
            self.ctx.get_new_instr_id(), 
            Some(vreg.clone()));


        let i = self.mir_arena.add_instr(instr);
        self.instr2reg.insert(i, vreg);
        // self.add_instr2cur_bb(i);
        i
    }
}
